package com.pine.app.module.security.oauth.provider.token.store;

import com.pine.app.module.security.core.common.AuthConstant;
import com.pine.app.module.security.oauth.provider.token.OAuth2AccessToken;
import com.pine.app.module.security.oauth.provider.OAuth2Authentication;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.dao.DataAccessException;
import org.springframework.data.redis.connection.RedisConnection;
import org.springframework.data.redis.core.RedisCallback;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.serializer.RedisSerializer;
import org.springframework.lang.Nullable;

/**
 * @author xiaoyuan
 * @create 2019/10/23 16:07
 **/
public class RedisTokenStore extends AbstractTokenStore {

    private  static final  String ACCESS_TOKEN= AuthConstant.CACHE_NAMESPACE+"TOKEN:";
    private  static final  String ACCESS_AUTH= AuthConstant.CACHE_NAMESPACE+"TOKEN:AUTH:";
    private  static final  String AUTH_TOKEN= AuthConstant.CACHE_NAMESPACE+"AUTH:TOKEN:";
    private  static final  String REFRESH_TOKEN= AuthConstant.CACHE_NAMESPACE+"REFRESH:";
    private  static final  String REFRESH_AUTH= AuthConstant.CACHE_NAMESPACE+"REFRESH:AUTH";
    private RedisTemplate   redisTemplate;

    @Autowired
    public void setRedisTemplate(RedisTemplate redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    public RedisTokenStore(){

    }

    public RedisTokenStore(RedisTemplate   redisTemplate){
        this.redisTemplate = redisTemplate;
    }
    @Override
    public OAuth2Authentication readAuthentication(OAuth2AccessToken token) {
        return readAuthentication(token.getAccessToken());
    }

    @Override
    public OAuth2Authentication readAuthentication(String key) {
        return (OAuth2Authentication) redisTemplate.opsForValue().get(ACCESS_AUTH+key);
    }

    @Override
    public void storeAccessToken(OAuth2AccessToken token, OAuth2Authentication auth2Authentication) {

        redisTemplate.executePipelined(new RedisCallback<Object>() {

            @Nullable
            @Override
            public Object doInRedis(RedisConnection redisConnection) throws DataAccessException {

                try{
                    redisConnection.openPipeline();
                    RedisSerializer<OAuth2AccessToken> auth2AccessTokenRedisSerializer =   redisTemplate.getValueSerializer();
                    RedisSerializer<OAuth2Authentication> oAuth2AuthenticationRedisSerializer =   redisTemplate.getValueSerializer();
                    RedisSerializer<String> keySerializer = (RedisSerializer<String>) redisTemplate.getKeySerializer();
                    byte[] access_token_key =  keySerializer.serialize(ACCESS_TOKEN+token.getAccessToken());
                    byte[] access_auth_key =  keySerializer.serialize(ACCESS_AUTH+token.getAccessToken());
                    byte[] auth_token_key =  keySerializer.serialize(AUTH_TOKEN+authenticationKeyGenerator.extractKey(auth2Authentication));
                    redisConnection.set(access_token_key,auth2AccessTokenRedisSerializer.serialize(token));
                    redisConnection.set(access_auth_key,oAuth2AuthenticationRedisSerializer.serialize(auth2Authentication));
                    redisConnection.set(auth_token_key,
                            auth2AccessTokenRedisSerializer.serialize(token));
                    if(token.getOauth2RefreshToken()!=null){
                        byte[] refresh_token_key =  keySerializer.serialize(REFRESH_TOKEN+token.getOauth2RefreshToken().getValue());
                        byte[] refresh_auth_key =  keySerializer.serialize(REFRESH_AUTH+token.getOauth2RefreshToken().getValue());
                        redisConnection.set(refresh_token_key,auth2AccessTokenRedisSerializer.serialize(token));
                        redisConnection.set(refresh_auth_key,oAuth2AuthenticationRedisSerializer.serialize(auth2Authentication));

                        redisConnection.expire(access_token_key,Long.valueOf((token.getOauth2RefreshToken().getExpiration().getTime()
                                - System.currentTimeMillis()) / 1000L).intValue());
                        redisConnection.expire(refresh_auth_key,Long.valueOf((token.getOauth2RefreshToken().getExpiration().getTime()
                                - System.currentTimeMillis()) / 1000L).intValue());
                    }
                    long date = Long.valueOf((token.getExpireTime().getTime()
                            - System.currentTimeMillis()) / 1000L).intValue();
                    redisConnection.expire(access_token_key, date);
                    redisConnection.expire(access_auth_key,date);
                    redisConnection.expire(auth_token_key,date);

                }finally {
                    redisConnection.closePipeline();
                    redisConnection.close();
                }
                return null;
            }
        });
    }

    @Override
    public OAuth2AccessToken readAccessToken(String token) {
        return (OAuth2AccessToken) redisTemplate.execute(new RedisCallback<OAuth2AccessToken>() {
            @Nullable
            @Override
            public OAuth2AccessToken doInRedis(RedisConnection redisConnection) throws DataAccessException {
                RedisSerializer<String> keySerializer = (RedisSerializer<String>) redisTemplate.getKeySerializer();
                RedisSerializer<OAuth2AccessToken> auth2AccessTokenRedisSerializer =   redisTemplate.getValueSerializer();
                byte[] auth_token_key =  keySerializer.serialize(ACCESS_TOKEN+token);
                return auth2AccessTokenRedisSerializer.deserialize(redisConnection.get(auth_token_key));
            }
        });
    }

    @Override
    public void removeAccessToken(OAuth2AccessToken auth2AccessToken) {
        removeAccessToken(auth2AccessToken.getAccessToken());
    }

    @Override
    public void removeAccessToken(String token) {

        OAuth2AccessToken accessToken =  this.readAccessToken(token);
        if(accessToken!=null){
            redisTemplate.delete(ACCESS_TOKEN+token);
            if(accessToken.getOauth2RefreshToken()!=null){
                this.redisTemplate.delete(REFRESH_TOKEN+accessToken.getRefreshToken());
                this.redisTemplate.delete(REFRESH_AUTH+accessToken.getRefreshToken());
            }
        }
        OAuth2Authentication oAuth2Authentication = this.readAuthentication(token);
        if (oAuth2Authentication != null) {
            redisTemplate.delete(AUTH_TOKEN+this.authenticationKeyGenerator.extractKey(oAuth2Authentication));
            redisTemplate.delete(ACCESS_AUTH+token);
        }
    }

    @Override
    public void removeRefreshToken(String token) {
        this.redisTemplate.delete(REFRESH_TOKEN+token);
        this.redisTemplate.delete(REFRESH_AUTH+token);
    }

    @Override
    public OAuth2AccessToken readRefreshToken(String refreshToken) {
        return (OAuth2AccessToken) redisTemplate.opsForValue().get(REFRESH_TOKEN+refreshToken);
    }

    @Override
    public OAuth2Authentication readRefreshToAuthentication(String refreshToken) {
        return (OAuth2Authentication) redisTemplate.opsForValue().get(REFRESH_AUTH+refreshToken);
    }

    @Override
    public OAuth2AccessToken getAccessToken(OAuth2Authentication auth2Authentication) {
        return (OAuth2AccessToken) redisTemplate.opsForValue().get(AUTH_TOKEN+authenticationKeyGenerator.extractKey(auth2Authentication));
    }
}
